# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false -tf-set-original-tf-func-name-on-import=true %s -o - | FileCheck %s

node {
  name: "custom_relu_func_call"
  op: "custom_relu"
}
node {
  name: "custom_embedding_matmul_func_call"
  op: "custom_embedding_matmul"
}
library {
  function {
    signature {
      name: "custom_relu"
    }
    attr {
      key: "_implements"
      value {
        func {
          name: "tensorflow.relu"
        }
      }
    }
  }
  function {
    signature {
      name: "custom_embedding_matmul"
    }
    attr {
      key: "_implements"
      value {
        func {
          name: "tensorflow.embedding_matmul"
          attr {
            key: "key1"
            value {
              i: 2
            }
          }
          attr {
            key: "key2"
            value {
              b: false
            }
          }
        }
      }
    }
  }
}

# CHECK-DAG: func private @custom_relu{{[0-9]*}}(){{.+}}tf._implements = #tf_type.func<@tensorflow.relu, {}>, tf._original_func_name = "custom_relu"}
# CHECK-DAG: func private @custom_embedding_matmul{{[0-9]*}}(){{.+}}tf._implements = #tf_type.func<@tensorflow.embedding_matmul, {key1 = 2 : i64, key2 = false}>, tf._original_func_name = "custom_embedding_matmul"}
